import sys
import os

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
from image_processing_modules.clip_model import ChineseClipModelWrapper

import numpy as np
from configs.config import Config


def compute_similarity(
    embedding1: np.ndarray, embedding2: np.ndarray, metric: str = "cosine"
) -> float:
    from utils.numpy_utils import calculate_similarity

    # 计算原始相似度/距离
    return calculate_similarity(embedding1, embedding2, metric=metric)


# ---------------------- 测试用例 ----------------------
if __name__ == "__main__":
    # 初始化中文CLIP模型（使用作者推荐的模型）
    
    # 加载配置
    config_path = "configs/local_22_config.json"
    config = Config(config_path)
    
    # 获取CLIP配置
    model_name = config.get("clip.model_name", "OFA-Sys/chinese-clip-vit-base-patch16")
    local_model_path = config.get("clip.local_model_path", "")
    offline_mode = config.get("clip.offline_mode", True)
    device = config.get("clip.device", "cpu")  # 使用CPU避免设备问题
    
    #chinese_clip = ChineseClipModelWrapper("OFA-Sys/chinese-clip-vit-large-patch14")
    print("正在初始化Chinese-CLIP模型...")
    chinese_clip = ChineseClipModelWrapper(
        model_name=model_name,
        local_model_path=local_model_path,
        offline_mode=offline_mode,
        device=device
    )
    
    print("Chinese-CLIP模型初始化成功!")
    print(f"模型嵌入维度: {chinese_clip.embedding_dim}")

    # 测试1：图像向量提取
    test_image_path = "D:/Temp/Pic/test/test_04.png"  # 替换为真实熊猫图像路径
    try:
        image_emb = chinese_clip.image_to_embedding(test_image_path)
        print(
            f"测试1：图像向量形状：{image_emb.shape}，L2范数：{np.linalg.norm(image_emb):.4f}，示例值：{image_emb[:15]}..."
        )  # 应输出(1024,)，范数≈1.0
    except Exception as e:
        print(f"测试1失败：{e}")

    # 测试2：文本向量生成（中文）
    test_texts: list[str] = ["熊猫", "熊", "猫", "北极熊", "鲸", "孔雀"]
    try:
        print("测试2：文本向量生成和图片的比对")
        idx = 1
        for test_text in test_texts:
            text_emb = chinese_clip.text_to_embedding(test_text)
            print(
                f"{idx}：文本：{test_text}，文本向量形状：{text_emb.shape}，L2范数：{np.linalg.norm(text_emb):.4f}"
            )  # 应输出(1024,)，范数≈1.0
            sim: float = compute_similarity(image_emb, text_emb)
            print(f"{idx}：图像与文本的相似度相似度：{sim}")  # 预期提升至0.5+
            idx += 1

    except Exception as e:
        print(f"测试2失败：{e}")

    # 测试3：批量图像向量提取
    test_image_paths = [
        "D:/Temp/Pic/test/test_04.png",
        "D:/Temp/Pic/test/test_03.png",
    ]  # 替换为真实图像路径
    try:
        batch_emb = chinese_clip.batch_image_to_embeddings(test_image_paths)
        print(
            f"测试3：批量向量形状：{batch_emb.shape}，每行范数≈1.0：{np.allclose(np.linalg.norm(batch_emb, axis=1), 1.0, atol=1e-6)}"
        )  # 应输出(2, 1024)，True
    except Exception as e:
        print(f"测试3失败：{e}")
